package com.swak.loadbalance.impl;

import java.util.List;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.atomic.AtomicLong;

import com.swak.loadbalance.Referer;

/**
 * 加权轮询
 *
 * @author lifeng
 * @date 2020年4月30日 上午10:25:39
 */
public class WeightingRoundRobinLoadBalance<T> extends AbstractLoadBalance<T> {

	private static final int RECYCLE_PERIOD = 60000;

	private ConcurrentMap<Integer, WeightedRoundRobin> weightMap = new ConcurrentHashMap<>();

	@Override
	public void onRefresh(List<T> referers) {
		super.onRefresh(referers);
		this.weightMap.clear();
	}

	@Override
	protected Referer<T> doSelect() {
		List<Referer<T>> referers = this.getReferers();
		int totalWeight = 0;
		long maxCurrent = Long.MIN_VALUE;
		long now = System.currentTimeMillis();
		Referer<T> selectedInvoker = null;
		WeightedRoundRobin selectedWRR = null;
		for (Referer<T> invoker : referers) {
			int identify = invoker.hashCode();
			int weight = invoker.getWeight();
			WeightedRoundRobin weightedRoundRobin = weightMap.computeIfAbsent(identify, k -> {
				WeightedRoundRobin wrr = new WeightedRoundRobin();
				wrr.setWeight(weight);
				return wrr;
			});

			if (weight != weightedRoundRobin.getWeight()) {
				// weight changed
				weightedRoundRobin.setWeight(weight);
			}
			long cur = weightedRoundRobin.increaseCurrent();
			weightedRoundRobin.setLastUpdate(now);
			if (cur > maxCurrent) {
				maxCurrent = cur;
				selectedInvoker = invoker;
				selectedWRR = weightedRoundRobin;
			}
			totalWeight += weight;
		}
		if (referers.size() != weightMap.size()) {
			weightMap.entrySet().removeIf(item -> now - item.getValue().getLastUpdate() > RECYCLE_PERIOD);
		}
		if (selectedInvoker != null) {
			selectedWRR.sel(totalWeight);
			return selectedInvoker;
		}
		// should not happen here
		return referers.get(0);
	}

	protected static class WeightedRoundRobin {
		private int weight;
		private AtomicLong current = new AtomicLong(0);
		private long lastUpdate;

		public int getWeight() {
			return weight;
		}

		public void setWeight(int weight) {
			this.weight = weight;
			current.set(0);
		}

		public long increaseCurrent() {
			return current.addAndGet(weight);
		}

		public void sel(int total) {
			current.addAndGet(-1 * total);
		}

		public long getLastUpdate() {
			return lastUpdate;
		}

		public void setLastUpdate(long lastUpdate) {
			this.lastUpdate = lastUpdate;
		}
	}
}
